{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial 9: Recurrent GNNs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this tutorial we will implement an approximation of the Graph Neural Network Model (without enforcing contraction map) and analyze the GatedGraph Convolution of Pytorch Geometric."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os.path as osp\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch_geometric.transforms as T\n",
    "import torch_geometric\n",
    "from torch_geometric.datasets import Planetoid, TUDataset\n",
    "from torch_geometric.data import DataLoader\n",
    "from torch_geometric.nn.inits import uniform\n",
    "from torch.nn import Parameter as Param\n",
    "from torch import Tensor \n",
    "torch.manual_seed(42)\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "from torch_geometric.nn.conv import MessagePassing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = 'Cora'\n",
    "transform = T.Compose([\n",
    "    T.AddTrainValTestMask('train_rest', num_val=500, num_test=500),\n",
    "    T.TargetIndegree(),\n",
    "])\n",
    "path = osp.join('data', dataset)\n",
    "dataset = Planetoid(path, dataset, transform=transform)\n",
    "data = dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = 'Cora'\n",
    "path = osp.join('data', dataset)\n",
    "dataset = Planetoid(path, dataset, transform=T.NormalizeFeatures())\n",
    "data = dataset[0]\n",
    "data = data.to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Graph Neural Network Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"./transition.png\" width=\"500\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"./output.png\" width=\"500\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The MLP class is used to instantiate the transition and output functions as simple feed forard networks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MLP(nn.Module):\n",
    "    def __init__(self, input_dim, hid_dims, out_dim):\n",
    "        super(MLP, self).__init__()\n",
    "\n",
    "        self.mlp = nn.Sequential()\n",
    "        dims = [input_dim] + hid_dims + [out_dim]\n",
    "        for i in range(len(dims)-1):\n",
    "            self.mlp.add_module('lay_{}'.format(i),nn.Linear(in_features=dims[i], out_features=dims[i+1]))\n",
    "            if i+2 < len(dims):\n",
    "                self.mlp.add_module('act_{}'.format(i), nn.Tanh())\n",
    "    def reset_parameters(self):\n",
    "        for i, l in enumerate(self.mlp):\n",
    "            if type(l) == nn.Linear:\n",
    "                nn.init.xavier_normal_(l.weight)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.mlp(x)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The GNNM calss puts together the state propagations and the readout of the nodes' states."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GNNM(MessagePassing):\n",
    "    def __init__(self, n_nodes, out_channels, features_dim, hid_dims, num_layers = 50, eps=1e-3, aggr = 'add',\n",
    "                 bias = True, **kwargs):\n",
    "        super(GNNM, self).__init__(aggr=aggr, **kwargs)\n",
    "\n",
    "        self.node_states = Param(torch.zeros((n_nodes, features_dim)), requires_grad=False)\n",
    "        self.out_channels = out_channels\n",
    "        self.eps = eps\n",
    "        self.num_layers = num_layers\n",
    "        \n",
    "        self.transition = MLP(features_dim, hid_dims, features_dim)\n",
    "        self.readout = MLP(features_dim, hid_dims, out_channels)\n",
    "        \n",
    "        self.reset_parameters()\n",
    "        print(self.transition)\n",
    "        print(self.readout)\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        self.transition.reset_parameters()\n",
    "        self.readout.reset_parameters()\n",
    "        \n",
    "    def forward(self): \n",
    "        edge_index = data.edge_index\n",
    "        edge_weight = data.edge_attr\n",
    "        node_states = self.node_states\n",
    "        for i in range(self.num_layers):\n",
    "            m = self.propagate(edge_index, x=node_states, edge_weight=edge_weight,\n",
    "                               size=None)\n",
    "            new_states = self.transition(m)\n",
    "            with torch.no_grad():\n",
    "                distance = torch.norm(new_states - node_states, dim=1)\n",
    "                convergence = distance < self.eps\n",
    "            node_states = new_states\n",
    "            if convergence.all():\n",
    "                break\n",
    "            \n",
    "        out = self.readout(node_states)\n",
    "        \n",
    "        return F.log_softmax(out, dim=-1)\n",
    "\n",
    "    def message(self, x_j, edge_weight):\n",
    "        return x_j if edge_weight is None else edge_weight.view(-1, 1) * x_j\n",
    "\n",
    "    def message_and_aggregate(self, adj_t, x) :\n",
    "        return matmul(adj_t, x, reduce=self.aggr)\n",
    "\n",
    "    def __repr__(self):\n",
    "        return '{}({}, num_layers={})'.format(self.__class__.__name__,\n",
    "                                              self.out_channels,\n",
    "                                              self.num_layers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MLP(\n",
      "  (mlp): Sequential(\n",
      "    (lay_0): Linear(in_features=32, out_features=64, bias=True)\n",
      "    (act_0): Tanh()\n",
      "    (lay_1): Linear(in_features=64, out_features=64, bias=True)\n",
      "    (act_1): Tanh()\n",
      "    (lay_2): Linear(in_features=64, out_features=64, bias=True)\n",
      "    (act_2): Tanh()\n",
      "    (lay_3): Linear(in_features=64, out_features=64, bias=True)\n",
      "    (act_3): Tanh()\n",
      "    (lay_4): Linear(in_features=64, out_features=64, bias=True)\n",
      "    (act_4): Tanh()\n",
      "    (lay_5): Linear(in_features=64, out_features=32, bias=True)\n",
      "  )\n",
      ")\n",
      "MLP(\n",
      "  (mlp): Sequential(\n",
      "    (lay_0): Linear(in_features=32, out_features=64, bias=True)\n",
      "    (act_0): Tanh()\n",
      "    (lay_1): Linear(in_features=64, out_features=64, bias=True)\n",
      "    (act_1): Tanh()\n",
      "    (lay_2): Linear(in_features=64, out_features=64, bias=True)\n",
      "    (act_2): Tanh()\n",
      "    (lay_3): Linear(in_features=64, out_features=64, bias=True)\n",
      "    (act_3): Tanh()\n",
      "    (lay_4): Linear(in_features=64, out_features=64, bias=True)\n",
      "    (act_4): Tanh()\n",
      "    (lay_5): Linear(in_features=64, out_features=7, bias=True)\n",
      "  )\n",
      ")\n",
      "Epoch: 001, Train Acc: 0.12857, Val Acc: 0.06800, Test Acc: 0.08800\n",
      "Epoch: 002, Train Acc: 0.14286, Val Acc: 0.25200, Test Acc: 0.25100\n",
      "Epoch: 003, Train Acc: 0.12143, Val Acc: 0.24400, Test Acc: 0.26100\n",
      "Epoch: 004, Train Acc: 0.17143, Val Acc: 0.20200, Test Acc: 0.20000\n",
      "Epoch: 005, Train Acc: 0.15714, Val Acc: 0.23000, Test Acc: 0.22700\n",
      "Epoch: 006, Train Acc: 0.21429, Val Acc: 0.09800, Test Acc: 0.10100\n",
      "Epoch: 007, Train Acc: 0.14286, Val Acc: 0.11600, Test Acc: 0.09800\n",
      "Epoch: 008, Train Acc: 0.17857, Val Acc: 0.16400, Test Acc: 0.14800\n",
      "Epoch: 009, Train Acc: 0.13571, Val Acc: 0.10200, Test Acc: 0.10200\n",
      "Epoch: 010, Train Acc: 0.13571, Val Acc: 0.10800, Test Acc: 0.10600\n",
      "Epoch: 011, Train Acc: 0.15000, Val Acc: 0.09800, Test Acc: 0.10000\n",
      "Epoch: 012, Train Acc: 0.17143, Val Acc: 0.11200, Test Acc: 0.11000\n",
      "Epoch: 013, Train Acc: 0.17857, Val Acc: 0.10400, Test Acc: 0.09800\n",
      "Epoch: 014, Train Acc: 0.18571, Val Acc: 0.13000, Test Acc: 0.10500\n",
      "Epoch: 015, Train Acc: 0.18571, Val Acc: 0.20600, Test Acc: 0.17300\n",
      "Epoch: 016, Train Acc: 0.20714, Val Acc: 0.18000, Test Acc: 0.16700\n",
      "Epoch: 017, Train Acc: 0.19286, Val Acc: 0.13600, Test Acc: 0.14100\n",
      "Epoch: 018, Train Acc: 0.17857, Val Acc: 0.11000, Test Acc: 0.11400\n",
      "Epoch: 019, Train Acc: 0.18571, Val Acc: 0.11800, Test Acc: 0.12000\n",
      "Epoch: 020, Train Acc: 0.18571, Val Acc: 0.13200, Test Acc: 0.13200\n",
      "Epoch: 021, Train Acc: 0.20000, Val Acc: 0.17400, Test Acc: 0.15900\n",
      "Epoch: 022, Train Acc: 0.19286, Val Acc: 0.20800, Test Acc: 0.16800\n",
      "Epoch: 023, Train Acc: 0.13571, Val Acc: 0.15000, Test Acc: 0.12300\n",
      "Epoch: 024, Train Acc: 0.15714, Val Acc: 0.10200, Test Acc: 0.09300\n",
      "Epoch: 025, Train Acc: 0.14286, Val Acc: 0.08200, Test Acc: 0.08800\n",
      "Epoch: 026, Train Acc: 0.12857, Val Acc: 0.12400, Test Acc: 0.11700\n",
      "Epoch: 027, Train Acc: 0.15000, Val Acc: 0.11800, Test Acc: 0.11200\n",
      "Epoch: 028, Train Acc: 0.16429, Val Acc: 0.18000, Test Acc: 0.16500\n",
      "Epoch: 029, Train Acc: 0.14286, Val Acc: 0.15600, Test Acc: 0.14500\n",
      "Epoch: 030, Train Acc: 0.15000, Val Acc: 0.15800, Test Acc: 0.14300\n",
      "Epoch: 031, Train Acc: 0.18571, Val Acc: 0.17200, Test Acc: 0.15200\n",
      "Epoch: 032, Train Acc: 0.21429, Val Acc: 0.16000, Test Acc: 0.16800\n",
      "Epoch: 033, Train Acc: 0.22857, Val Acc: 0.24200, Test Acc: 0.27200\n",
      "Epoch: 034, Train Acc: 0.25000, Val Acc: 0.21200, Test Acc: 0.23800\n",
      "Epoch: 035, Train Acc: 0.22857, Val Acc: 0.17200, Test Acc: 0.17500\n",
      "Epoch: 036, Train Acc: 0.25000, Val Acc: 0.13800, Test Acc: 0.15700\n",
      "Epoch: 037, Train Acc: 0.25000, Val Acc: 0.13800, Test Acc: 0.14500\n",
      "Epoch: 038, Train Acc: 0.23571, Val Acc: 0.14200, Test Acc: 0.14800\n",
      "Epoch: 039, Train Acc: 0.28571, Val Acc: 0.15400, Test Acc: 0.15400\n",
      "Epoch: 040, Train Acc: 0.15000, Val Acc: 0.05800, Test Acc: 0.06500\n",
      "Epoch: 041, Train Acc: 0.13571, Val Acc: 0.06400, Test Acc: 0.08300\n",
      "Epoch: 042, Train Acc: 0.08571, Val Acc: 0.05000, Test Acc: 0.05700\n",
      "Epoch: 043, Train Acc: 0.13571, Val Acc: 0.06200, Test Acc: 0.06400\n",
      "Epoch: 044, Train Acc: 0.14286, Val Acc: 0.16000, Test Acc: 0.14200\n",
      "Epoch: 045, Train Acc: 0.14286, Val Acc: 0.16000, Test Acc: 0.14200\n",
      "Epoch: 046, Train Acc: 0.14286, Val Acc: 0.16000, Test Acc: 0.14200\n",
      "Epoch: 047, Train Acc: 0.06429, Val Acc: 0.08800, Test Acc: 0.06500\n",
      "Epoch: 048, Train Acc: 0.12857, Val Acc: 0.06600, Test Acc: 0.06500\n",
      "Epoch: 049, Train Acc: 0.13571, Val Acc: 0.06600, Test Acc: 0.06700\n",
      "Epoch: 050, Train Acc: 0.15000, Val Acc: 0.16000, Test Acc: 0.14600\n"
     ]
    }
   ],
   "source": [
    "model = GNNM(data.num_nodes, dataset.num_classes, 32, [64,64,64,64,64], eps=0.01).to(device)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n",
    "loss_fn = nn.CrossEntropyLoss()\n",
    "\n",
    "\n",
    "test_dataset = dataset[:len(dataset) // 10]\n",
    "train_dataset = dataset[len(dataset) // 10:]\n",
    "test_loader = DataLoader(test_dataset)\n",
    "train_loader = DataLoader(train_dataset)\n",
    "\n",
    "def train():\n",
    "    model.train()\n",
    "    optimizer.zero_grad()\n",
    "    loss_fn(model()[data.train_mask], data.y[data.train_mask]).backward()\n",
    "    optimizer.step()\n",
    "\n",
    "\n",
    "def test():\n",
    "    model.eval()\n",
    "    logits, accs = model(), []\n",
    "    for _, mask in data('train_mask', 'val_mask', 'test_mask'):\n",
    "        pred = logits[mask].max(1)[1]\n",
    "        acc = pred.eq(data.y[mask]).sum().item() / mask.sum().item()\n",
    "        accs.append(acc)\n",
    "    return accs\n",
    "\n",
    "\n",
    "for epoch in range(1, 51):\n",
    "    train()\n",
    "    accs = test()\n",
    "    train_acc = accs[0]\n",
    "    val_acc = accs[1]\n",
    "    test_acc = accs[2]\n",
    "    print('Epoch: {:03d}, Train Acc: {:.5f}, '\n",
    "          'Val Acc: {:.5f}, Test Acc: {:.5f}'.format(epoch, train_acc,\n",
    "                                                       val_acc, test_acc))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Gated Graph Neural Network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GatedGraphConv(MessagePassing):\n",
    "    \n",
    "    def __init__(self, out_channels, num_layers, aggr = 'add',\n",
    "                 bias = True, **kwargs):\n",
    "        super(GatedGraphConv, self).__init__(aggr=aggr, **kwargs)\n",
    "\n",
    "        self.out_channels = out_channels\n",
    "        self.num_layers = num_layers\n",
    "\n",
    "        self.weight = Param(Tensor(num_layers, out_channels, out_channels))\n",
    "        self.rnn = torch.nn.GRUCell(out_channels, out_channels, bias=bias)\n",
    "\n",
    "        self.reset_parameters()\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        uniform(self.out_channels, self.weight)\n",
    "        self.rnn.reset_parameters()\n",
    "\n",
    "    def forward(self, data):\n",
    "        \"\"\"\"\"\"\n",
    "        x = data.x\n",
    "        edge_index = data.edge_index\n",
    "        edge_weight = data.edge_attr\n",
    "        if x.size(-1) > self.out_channels:\n",
    "            raise ValueError('The number of input channels is not allowed to '\n",
    "                             'be larger than the number of output channels')\n",
    "\n",
    "        if x.size(-1) < self.out_channels:\n",
    "            zero = x.new_zeros(x.size(0), self.out_channels - x.size(-1))\n",
    "            x = torch.cat([x, zero], dim=1)\n",
    "\n",
    "        for i in range(self.num_layers):\n",
    "            m = torch.matmul(x, self.weight[i])\n",
    "            m = self.propagate(edge_index, x=m, edge_weight=edge_weight,\n",
    "                               size=None)\n",
    "            x = self.rnn(m, x)\n",
    "\n",
    "        return x\n",
    "\n",
    "    def message(self, x_j, edge_weight):\n",
    "        return x_j if edge_weight is None else edge_weight.view(-1, 1) * x_j\n",
    "\n",
    "    def message_and_aggregate(self, adj_t, x):\n",
    "        return matmul(adj_t, x, reduce=self.aggr)\n",
    "\n",
    "    def __repr__(self):\n",
    "        return '{}({}, num_layers={})'.format(self.__class__.__name__,\n",
    "                                              self.out_channels,\n",
    "                                              self.num_layers)\n",
    "\n",
    "class GGNN(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super(GGNN, self).__init__()\n",
    "        \n",
    "        self.conv = GatedGraphConv(1433, 3)\n",
    "        self.mlp = MLP(1433, [32,32,32], dataset.num_classes)\n",
    "        \n",
    "    def forward(self):\n",
    "        x = self.conv(data)\n",
    "        x = self.mlp(x)\n",
    "        return F.log_softmax(x, dim=-1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 001, Train Acc: 0.20714, Val Acc: 0.18200, Test Acc: 0.17500\n",
      "Epoch: 002, Train Acc: 0.27143, Val Acc: 0.24000, Test Acc: 0.22200\n",
      "Epoch: 003, Train Acc: 0.24286, Val Acc: 0.23400, Test Acc: 0.22200\n",
      "Epoch: 004, Train Acc: 0.28571, Val Acc: 0.26200, Test Acc: 0.25400\n",
      "Epoch: 005, Train Acc: 0.40714, Val Acc: 0.35200, Test Acc: 0.33700\n",
      "Epoch: 006, Train Acc: 0.39286, Val Acc: 0.38200, Test Acc: 0.36300\n",
      "Epoch: 007, Train Acc: 0.39286, Val Acc: 0.36200, Test Acc: 0.34000\n",
      "Epoch: 008, Train Acc: 0.47857, Val Acc: 0.44200, Test Acc: 0.42700\n",
      "Epoch: 009, Train Acc: 0.44286, Val Acc: 0.43200, Test Acc: 0.40400\n",
      "Epoch: 010, Train Acc: 0.63571, Val Acc: 0.53600, Test Acc: 0.54500\n",
      "Epoch: 011, Train Acc: 0.67857, Val Acc: 0.55800, Test Acc: 0.56300\n",
      "Epoch: 012, Train Acc: 0.67857, Val Acc: 0.56200, Test Acc: 0.55900\n",
      "Epoch: 013, Train Acc: 0.71429, Val Acc: 0.58200, Test Acc: 0.57900\n",
      "Epoch: 014, Train Acc: 0.77143, Val Acc: 0.59800, Test Acc: 0.61700\n",
      "Epoch: 015, Train Acc: 0.79286, Val Acc: 0.61400, Test Acc: 0.62800\n",
      "Epoch: 016, Train Acc: 0.83571, Val Acc: 0.66800, Test Acc: 0.65400\n",
      "Epoch: 017, Train Acc: 0.85714, Val Acc: 0.66000, Test Acc: 0.65000\n",
      "Epoch: 018, Train Acc: 0.88571, Val Acc: 0.65200, Test Acc: 0.64600\n",
      "Epoch: 019, Train Acc: 0.93571, Val Acc: 0.64000, Test Acc: 0.63200\n",
      "Epoch: 020, Train Acc: 0.95000, Val Acc: 0.62000, Test Acc: 0.61900\n",
      "Epoch: 021, Train Acc: 0.94286, Val Acc: 0.59400, Test Acc: 0.59500\n",
      "Epoch: 022, Train Acc: 0.94286, Val Acc: 0.60000, Test Acc: 0.61200\n",
      "Epoch: 023, Train Acc: 0.93571, Val Acc: 0.62400, Test Acc: 0.62100\n",
      "Epoch: 024, Train Acc: 0.96429, Val Acc: 0.67600, Test Acc: 0.66900\n",
      "Epoch: 025, Train Acc: 0.95714, Val Acc: 0.64800, Test Acc: 0.63900\n",
      "Epoch: 026, Train Acc: 0.96429, Val Acc: 0.64800, Test Acc: 0.62500\n",
      "Epoch: 027, Train Acc: 0.93571, Val Acc: 0.65400, Test Acc: 0.62200\n",
      "Epoch: 028, Train Acc: 0.97857, Val Acc: 0.66000, Test Acc: 0.63500\n",
      "Epoch: 029, Train Acc: 0.96429, Val Acc: 0.65800, Test Acc: 0.64400\n",
      "Epoch: 030, Train Acc: 0.96429, Val Acc: 0.66400, Test Acc: 0.65200\n",
      "Epoch: 031, Train Acc: 0.97857, Val Acc: 0.65800, Test Acc: 0.66100\n",
      "Epoch: 032, Train Acc: 0.97857, Val Acc: 0.66600, Test Acc: 0.65600\n",
      "Epoch: 033, Train Acc: 0.97857, Val Acc: 0.65000, Test Acc: 0.63900\n",
      "Epoch: 034, Train Acc: 0.97857, Val Acc: 0.64200, Test Acc: 0.62900\n",
      "Epoch: 035, Train Acc: 0.97857, Val Acc: 0.63800, Test Acc: 0.62800\n",
      "Epoch: 036, Train Acc: 0.97857, Val Acc: 0.65200, Test Acc: 0.63100\n",
      "Epoch: 037, Train Acc: 0.97857, Val Acc: 0.66400, Test Acc: 0.63600\n",
      "Epoch: 038, Train Acc: 0.97857, Val Acc: 0.67000, Test Acc: 0.64000\n",
      "Epoch: 039, Train Acc: 0.97857, Val Acc: 0.67000, Test Acc: 0.64000\n",
      "Epoch: 040, Train Acc: 0.97857, Val Acc: 0.66800, Test Acc: 0.64100\n",
      "Epoch: 041, Train Acc: 0.97857, Val Acc: 0.66800, Test Acc: 0.64100\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-12-fde0cb8923bf>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     28\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mepoch\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m51\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     29\u001b[0m     \u001b[0mtrain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 30\u001b[0;31m     \u001b[0maccs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtest\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     31\u001b[0m     \u001b[0mtrain_acc\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0maccs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     32\u001b[0m     \u001b[0mval_acc\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0maccs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-12-fde0cb8923bf>\u001b[0m in \u001b[0;36mtest\u001b[0;34m()\u001b[0m\n\u001b[1;32m     20\u001b[0m     \u001b[0mlogits\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0maccs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     21\u001b[0m     \u001b[0;32mfor\u001b[0m \u001b[0m_\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmask\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mdata\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'train_mask'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'val_mask'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'test_mask'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 22\u001b[0;31m         \u001b[0mpred\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlogits\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mmask\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmax\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     23\u001b[0m         \u001b[0macc\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpred\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0meq\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdata\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0my\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mmask\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msum\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mitem\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m/\u001b[0m \u001b[0mmask\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msum\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mitem\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     24\u001b[0m         \u001b[0maccs\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0macc\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "model = GGNN().to(device)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n",
    "loss_fn = nn.CrossEntropyLoss()\n",
    "\n",
    "\n",
    "test_dataset = dataset[:len(dataset) // 10]\n",
    "train_dataset = dataset[len(dataset) // 10:]\n",
    "test_loader = DataLoader(test_dataset)\n",
    "train_loader = DataLoader(train_dataset)\n",
    "\n",
    "def train():\n",
    "    model.train()\n",
    "    optimizer.zero_grad()\n",
    "    loss_fn(model()[data.train_mask], data.y[data.train_mask]).backward()\n",
    "    optimizer.step()\n",
    "\n",
    "\n",
    "def test():\n",
    "    model.eval()\n",
    "    logits, accs = model(), []\n",
    "    for _, mask in data('train_mask', 'val_mask', 'test_mask'):\n",
    "        pred = logits[mask].max(1)[1]\n",
    "        acc = pred.eq(data.y[mask]).sum().item() / mask.sum().item()\n",
    "        accs.append(acc)\n",
    "    return accs\n",
    "\n",
    "\n",
    "for epoch in range(1, 51):\n",
    "    train()\n",
    "    accs = test()\n",
    "    train_acc = accs[0]\n",
    "    val_acc = accs[1]\n",
    "    test_acc = accs[2]\n",
    "    print('Epoch: {:03d}, Train Acc: {:.5f}, '\n",
    "          'Val Acc: {:.5f}, Test Acc: {:.5f}'.format(epoch, train_acc,\n",
    "                                                       val_acc, test_acc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
